介绍¶

Numba 是 python 的即时(Just-in-time)编译器,即当调用 python 函数时,全部或部分代码就会被转换为“即时”执行的机器码,它将以本地机器码速度运行!

在 Numba 的帮助下,可以加速所有计算负载比较大的 python 函数(例如循环)。它还支持 numpy 库。所以,也可以在计算中使用 numpy,并加快整体计算(因为 python 中的循环非常慢)。

使用Numba根本不需要为了获得一些的加速来改变Python的代码。只需要添加一个熟悉的 python 功能,即添加一个包装器(一个装饰器)到你的函数上。

例如:

from numba import jit
@jit
def function(x):
    # your loop or numerically intensive computations
    return x

如何使用 Numba?¶

Numba 使用 LLVM 编译器基础结构 将原生 python 代码转换成优化的机器码。使用 numba 运行代码的速度可与 C/C++ 或 Fortran 中的类似代码相媲美。

首先,Python 函数被传入,优化并转换为 numba 的中间表达,然后在类型推断(type inference)之后,就像 numpy 的类型推断(所以 python float 是一个 float64),它被转换为 LLVM 可解释代码。 然后将此代码提供给 LLVM 的即时编译器以生成机器码。

可以根据需要在运行时或导入时生成机器码,导入需要在 CPU(默认)或 GPU 上进行。

为了获得最佳性能,numba 实际上建议在 jit 装饰器中加上 nopython=True 参数,加上后就不会使用 Python 解释器了。或者也可以使用 @njit。如果您加上 nopython=True的装饰器失败并报错,可以用简单的 @jit 装饰器来编译部分代码,对于它能够编译的代码,将它们转换为函数,并编译成机器码。然后将其余部分代码提供给 python 解释器。

所以,只需要这样做:

from numba import njit, jit
@njit      # or @jit(nopython=True)
def function(a, b):
    # your loop or numerically intensive computations
    return result

当使用 @jit 时,请确保代码有 numba 可以编译的内容,比如包含库(numpy)和它支持的函数的计算密集型循环。否则它将不会编译任何东西,并且代码将比没有使用 numba 时更慢,因为存在 numba 内部代码检查的额外开销。

numba 会对首次作为机器码使用后的函数进行缓存。 因此,在第一次使用之后它将更快,因为它不需要再次编译这些代码,如果您使用的是和之前相同的参数类型。

如果代码是可并行化的,也可以传递 parallel=True 作为参数,但它必须与 nopython=True 一起使用,目前这只适用于CPU。

还可以指定希望函数具有的函数签名,但是这样就不会对您提供的任何其他类型的参数进行编译。 例如:

from numba import jit, int32
@jit(int32(int32, int32))
def function(a, b):
    # your loop or numerically intensive computations
    return result
# or if you haven't imported type names
# you can pass them as string
@jit('int32(int32, int32)')
def function(a, b):
    # your loop or numerically intensive computations
    return result

现在函数只能接收两个 int32 类型的参数并返回一个 int32 类型的值。 通过这种方式,可以更好地控制您的函数。 如果需要,您甚至可以传递多个函数签名。

还可以使用 numba 提供的其他装饰器:

  • @vectorize:允许将标量参数作为 numpy 的 ufuncs 使用,
  • @guvectorize:生成 NumPy 广义上的 ufuncs,
  • @stencil:定义一个函数使其成为 stencil 类型操作的核函数
  • @jitclass:用于 jit 类,
  • @cfunc:声明一个函数用于本地回调(被C/C++等调用),
  • @overload:注册您自己的函数实现,以便在 nopython 模式下使用,例如: @overload(scipy.special.j0)。

Numba 还有 Ahead of time(AOT)编译,它生成不依赖于 Numba 的已编译扩展模块。 但它只允许常规函数(ufuncs 就不行);必须指定函数签名。并且只能指定一种签名,如果需要指定多个签名,需要使用不同的名字。

它还根据您的CPU架构系列生成通用代码。

@vectorize 装饰器¶

通过使用 @vectorize 装饰器,可以对仅能对标量操作的函数进行转换,例如,如果您使用的是仅适用于标量的 python 的 math 库,则转换后就可以用于数组。 这提供了类似于 numpy 数组运算(ufuncs)的速度。 例如:

@vectorize
def func(a, b):
    # Some operation on scalars
    return result

还可以将 target 参数传递给此装饰器,该装饰器使 target 参数为 parallel 时用于并行化代码,为 cuda 时用于在 cuda\GPU 上运行代码。

@vectorize(target="parallel")
def func(a, b):
    # Some operation on scalars
    return result

使 target=“parallel” 或 “cuda” 进行矢量化通常比 numpy 实现的代码运行得更快,只要您的代码具有足够的计算密度或者数组足够大。如果不是,那么由于创建线程以及将元素分配到不同线程需要额外的开销,因此可能耗时更长。所以运算量应该足够大,才能获得明显的加速。

参考资料

  • https://mp.weixin.qq.com/s/VQl4XmXBN_CHUM-fpjjk1w

In [1]:
from numba import jit
@jit
def function(x):
    # your loop or numerically intensive computations
    return x